{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08717680-8435-4032-ac79-85964f1f45ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!/usr/bin/python\n",
    "# -*- coding: UTF-8 -*-\n",
    "# author : Liu Kun\n",
    "# date   : 2021-09-01 00:00:01\n",
    "\n",
    "import os\n",
    "import cmd\n",
    "import json\n",
    "import requests\n",
    "from bs4 import BeautifulSoup\n",
    "\n",
    "\n",
    "def request_paper_list(url, debug=False):\n",
    "    \"\"\"\n",
    "    根据链接爬取论文信息\n",
    "    :param url:   目标链接\n",
    "    :param debug: 调试信息打印开关[默认关闭]\n",
    "    :return: 论文信息列表\n",
    "    \"\"\"\n",
    "    # 请求页面\n",
    "    base_url = f\"https://openaccess.thecvf.com\"\n",
    "    response = requests.get(url)\n",
    "    # 使用BeautifulSoup加载页面\n",
    "    soup = BeautifulSoup(response.content, features=\"lxml\")\n",
    "    # 按标签查找\n",
    "    dl = soup.find(\"dl\")\n",
    "    dt_list = dl.find_all(\"dt\")\n",
    "    dd_list = dl.find_all(\"dd\")\n",
    "    # 去除Back\n",
    "    if \"back\" in dd_list[0].text.strip().lower():\n",
    "        dd_list.pop(0)\n",
    "    # 按奇偶处理dd标签\n",
    "    dd_odd_list = dd_list[::2]\n",
    "    dd_even_list = dd_list[1::2]\n",
    "    # 用于存放论文信息\n",
    "    papers = list()\n",
    "    # 逐行处理\n",
    "    for item in zip(dt_list, dd_odd_list, dd_even_list):\n",
    "        # Line 1\n",
    "        line_1 = item[0]\n",
    "        tag_a_1 = line_1.find(\"a\")\n",
    "        # Line 2\n",
    "        line_2 = item[1]\n",
    "        # Line 3\n",
    "        line_3 = item[2]\n",
    "        tag_a_2 = line_3.find(\"a\")\n",
    "        tag_div = line_3.find(\"div\")\n",
    "        paper_json = {\n",
    "            \"web_url\": base_url + tag_a_1[\"href\"],\n",
    "            \"title\": tag_a_1.text.strip(),\n",
    "            \"author\": line_2.text.strip().replace(\"\\n\", \"\").replace(\"\\r\", \"\"),\n",
    "            \"pdf_url\": base_url + tag_a_2[\"href\"],\n",
    "            \"infos\": tag_div.text.strip().replace(\"\\n\", \"\").replace(\"\\r\", \"\").replace(\"[bibtex]\", \"\"),\n",
    "        }\n",
    "        # 修复链接中缺少/链接符的问题\n",
    "        if not str(tag_a_1[\"href\"]).startswith(\"/\"):\n",
    "            paper_json[\"web_url\"] = base_url + \"/\" + tag_a_1[\"href\"]\n",
    "        if not str(tag_a_2[\"href\"]).startswith(\"/\"):\n",
    "            paper_json[\"pdf_url\"] = base_url + \"/\" + tag_a_2[\"href\"]\n",
    "        if debug:\n",
    "            print(\"-\" * 100)\n",
    "            print(json.dumps(paper_json, sort_keys=True, indent=4, separators=(',', ': ')))\n",
    "        papers.append(paper_json)\n",
    "    print(\"共爬取到论文信息%d条\" % len(papers))\n",
    "    return papers\n",
    "\n",
    "\n",
    "def export_jsonl(papers, filename):\n",
    "    \"\"\"\n",
    "    导出为jsonl格式\n",
    "    每行一条json数据\n",
    "    :param papers: 论文信息列表\n",
    "    :param filename: 导出文件名\n",
    "    :return: 无\n",
    "    \"\"\"\n",
    "    with open(filename, \"a+\") as jsonl:\n",
    "        for p in papers:\n",
    "            jsonl.write(json.dumps(p) + \"\\n\")\n",
    "    pass\n",
    "\n",
    "\n",
    "def export_url(papers, filename):\n",
    "    \"\"\"\n",
    "    导出为url格式\n",
    "    每行一条url数据\n",
    "    :param papers: 论文信息列表\n",
    "    :param filename: 导出文件名\n",
    "    :return: 无\n",
    "    \"\"\"\n",
    "    with open(filename, \"a+\") as txt:\n",
    "        for p in papers:\n",
    "            txt.write(p[\"pdf_url\"] + \"\\n\")\n",
    "    pass\n",
    "\n",
    "\n",
    "def downloader(url, folder):\n",
    "    \"\"\"\n",
    "    下载论文到指定目录\n",
    "    :param url: 论文链接\n",
    "    :param folder: 存储目录\n",
    "    :return: 本地文件名\n",
    "    \"\"\"\n",
    "    filename = url.split(\"/\")[-1]\n",
    "    file = os.path.join(folder, filename)\n",
    "    if os.path.exists(file):\n",
    "        print(f\"skip {file} @ {url}\")\n",
    "        return file\n",
    "    print(f\"download {file} from {url}\")\n",
    "    response = requests.get(url)\n",
    "    with open(file, \"wb\") as output:\n",
    "        output.write(response.content)\n",
    "    return file\n",
    "\n",
    "\n",
    "class CLI(cmd.Cmd):\n",
    "\n",
    "    intro = \"输入 help 或者 ? 查看帮助。\\n\"\n",
    "    prompt = \"thecvf>\"\n",
    "    __menu_list = None\n",
    "\n",
    "    def preloop(self):\n",
    "        with open(\"menu.json\", \"r\") as f:\n",
    "            self.__menu_list = json.load(f)\n",
    "\n",
    "    def do_crawler(self, _):\n",
    "        \"\"\"\n",
    "        爬取论文信息\n",
    "        :param _: 无需输入\n",
    "        :return: 无\n",
    "        \"\"\"\n",
    "        for target in self.__menu_list:\n",
    "            # 打印待爬取信息\n",
    "            print(target)\n",
    "            sub_path = target[\"path\"]\n",
    "            # 如文件夹不存在则创建\n",
    "            if not os.path.exists(sub_path):\n",
    "                os.makedirs(sub_path)\n",
    "            logfile = target[\"logfile\"]\n",
    "            target_url_list = target[\"links\"]\n",
    "            for target_url in target_url_list:\n",
    "                # 爬取论文链接\n",
    "                paper_list = request_paper_list(target_url)\n",
    "                # 记录论文信息到日志文件\n",
    "                fn = os.path.join(sub_path, logfile)\n",
    "                export_jsonl(paper_list, fn)\n",
    "                export_url(paper_list, fn.replace(\"jsonl\", \"txt\"))\n",
    "\n",
    "    def do_download(self, arg):\n",
    "        \"\"\"\n",
    "        下载论文\n",
    "        输入期望下载的目录名称，例如 CVPR2021\n",
    "        注意: 未输入任何名称时，将全部下载\n",
    "        :param arg: 即将下载的目录名称\n",
    "        :return: 无\n",
    "        \"\"\"\n",
    "        if arg is not None and len(arg) > 0:\n",
    "            for paper in self.__menu_list:\n",
    "                paper_path = paper[\"path\"]\n",
    "                if paper_path != arg:\n",
    "                    continue\n",
    "                paper_jsonl = paper[\"logfile\"]\n",
    "                print(f\"download {paper_path}\")\n",
    "                with open(f\"{paper_path}/{paper_jsonl}\", \"r\") as f:\n",
    "                    rows = f.readlines()\n",
    "                    for r in rows:\n",
    "                        obj = json.loads(r)\n",
    "                        downloader(obj[\"pdf_url\"], paper_path)\n",
    "        else:\n",
    "            for paper in self.__menu_list:\n",
    "                paper_path = paper[\"path\"]\n",
    "                paper_jsonl = paper[\"logfile\"]\n",
    "                print(f\"download {paper_path}\")\n",
    "                with open(f\"{paper_path}/{paper_jsonl}\", \"r\") as f:\n",
    "                    rows = f.readlines()\n",
    "                    for r in rows:\n",
    "                        obj = json.loads(r)\n",
    "                        downloader(obj[\"pdf_url\"], paper_path)\n",
    "        pass\n",
    "\n",
    "    def do_clean(self, arg):\n",
    "        \"\"\"\n",
    "        清理目录\n",
    "        输入期望清理的目录名称，例如 CVPR2021\n",
    "        注意: 未输入任何名称时，将全部清理\n",
    "        :param arg: 即将清理的目录名称\n",
    "        :return: 无\n",
    "        \"\"\"\n",
    "        if arg is not None and len(arg) > 0:\n",
    "            for paper in self.__menu_list:\n",
    "                path = paper[\"path\"]\n",
    "                if path != arg:\n",
    "                    continue\n",
    "                os.system(f\"rm -rf {path}\")\n",
    "                print(f\"remove {path}\")\n",
    "        else:\n",
    "            for paper in self.__menu_list:\n",
    "                path = paper[\"path\"]\n",
    "                os.system(f\"rm -rf {path}\")\n",
    "                print(f\"remove {path}\")\n",
    "        pass\n",
    "\n",
    "    def do_exit(self, _):\n",
    "        \"\"\"\n",
    "        退出命令行\n",
    "        :param _: 无需输入参数\n",
    "        :return: 无\n",
    "        \"\"\"\n",
    "        exit(0)\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    CLI().cmdloop()\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e31d24b-75e0-4f84-a500-458720c40253",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
